import argparse  # 用于解析命令行参数的库
from collections import namedtuple  # 提供轻量级的数据结构
from itertools import count  # 无限循环计数器

import os, sys, random  # 操作系统、系统操作及随机库
import numpy as np  # 用于数组操作和科学计算的库

import gym  # OpenAI Gym库，用于构建强化学习环境
import torch  # 深度学习框架 PyTorch
import torch.nn as nn  # 神经网络模块
import torch.nn.functional as F  # 提供激活函数和其他功能
import torch.optim as optim  # 优化器模块，用于梯度下降
from torch.distributions import Normal  # 正态分布，用于策略采样
from tensorboardX import SummaryWriter  # 用于记录训练日志

# 如果GPU可用，则使用CUDA，否则使用CPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# 创建一个命令行参数解析器
parser = argparse.ArgumentParser()

# 添加脚本运行时的参数
parser.add_argument('--mode', default='train', type=str)  # 模式：训练('train')或测试('test')
parser.add_argument("--env_name", default="Pendulum-v0")  # OpenAI Gym环境名称
parser.add_argument('--tau',  default=0.005, type=float)  # 目标网络的软更新系数
parser.add_argument('--target_update_interval', default=1, type=int)  # 目标网络更新间隔
parser.add_argument('--iteration', default=5, type=int)  # 迭代次数

# 学习相关参数
parser.add_argument('--learning_rate', default=3e-4, type=float)  # 学习率
parser.add_argument('--gamma', default=0.99, type=int)  # 折扣因子，用于奖励的衰减
parser.add_argument('--capacity', default=50000, type=int)  # 经验回放缓冲区大小
parser.add_argument('--num_iteration', default=100000, type=int)  # 总训练迭代次数
parser.add_argument('--batch_size', default=100, type=int)  # 批量大小
parser.add_argument('--seed', default=1, type=int)  # 随机种子，确保结果可复现

# 可选参数
parser.add_argument('--num_hidden_layers', default=2, type=int)  # 神经网络的隐藏层数
parser.add_argument('--sample_frequency', default=256, type=int)  # 采样频率
parser.add_argument('--activation', default='Relu', type=str)  # 激活函数类型
parser.add_argument('--render', default=False, type=bool)  # 是否显示渲染的界面
parser.add_argument('--log_interval', default=50, type=int)  # 日志记录间隔
parser.add_argument('--load', default=False, type=bool)  # 是否加载模型
parser.add_argument('--render_interval', default=100, type=int)  # 渲染间隔
parser.add_argument('--policy_noise', default=0.2, type=float)  # 策略噪声
parser.add_argument('--noise_clip', default=0.5, type=float)  # 噪声裁剪范围
parser.add_argument('--policy_delay', default=2, type=int)  # 策略更新延迟
parser.add_argument('--exploration_noise', default=0.1, type=float)  # 探索噪声
parser.add_argument('--max_episode', default=2000, type=int)  # 最大训练轮数
parser.add_argument('--print_log', default=5, type=int)  # 打印日志的间隔
args = parser.parse_args()  # 解析命令行参数

# 设置随机种子以保证结果可复现
# env.seed(args.seed)  # 环境随机种子
# torch.manual_seed(args.seed)  # PyTorch随机种子
# np.random.seed(args.seed)  # Numpy随机种子

# 获取当前脚本文件名
script_name = os.path.basename(__file__)

# 创建Gym环境
env = gym.make(args.env_name)

# 提取环境的状态空间和动作空间的维度
state_dim = env.observation_space.shape[0]  # 状态空间的维度
action_dim = env.action_space.shape[0]  # 动作空间的维度
max_action = float(env.action_space.high[0])  # 动作空间的最大值
min_Val = torch.tensor(1e-7).float().to(device)  # 防止数值错误的最小值

# 设置保存模型和日志的目录
directory = './exp' + script_name + args.env_name + './'


# 定义经验回放缓冲区类
class Replay_buffer():
    '''
    用于存储经验数据： (state, next_state, action, reward, done)
    '''
    def __init__(self, max_size=args.capacity):
        self.storage = []  # 存储经验的列表
        self.max_size = max_size  # 最大存储容量
        self.ptr = 0  # 指针，用于循环覆盖旧数据

    def push(self, data):
        # 如果缓冲区已满，则覆盖旧数据
        if len(self.storage) == self.max_size:
            self.storage[int(self.ptr)] = data  # 覆盖旧数据
            self.ptr = (self.ptr + 1) % self.max_size  # 更新指针位置
        else:
            self.storage.append(data)  # 添加新数据

    def sample(self, batch_size):
        # 随机从缓冲区中抽取一个批次的数据
        ind = np.random.randint(0, len(self.storage), size=batch_size)
        x, y, u, r, d = [], [], [], [], []  # 用于存储抽样结果

        # 遍历随机索引并提取对应的数据
        for i in ind:
            X, Y, U, R, D = self.storage[i]
            x.append(np.array(X, copy=False))
            y.append(np.array(Y, copy=False))
            u.append(np.array(U, copy=False))
            r.append(np.array(R, copy=False))
            d.append(np.array(D, copy=False))

        # 返回转换为NumPy数组的采样结果
        return np.array(x), np.array(y), np.array(u), np.array(r).reshape(-1, 1), np.array(d).reshape(-1, 1)


# 定义Actor（策略网络）类
class Actor(nn.Module):
    def __init__(self, state_dim, action_dim, max_action):
        super(Actor, self).__init__()
        # 定义三层全连接层，输入是状态，输出是动作
        self.fc1 = nn.Linear(state_dim, 400)  # 第一层全连接，400个神经元
        self.fc2 = nn.Linear(400, 300)  # 第二层全连接，300个神经元
        self.fc3 = nn.Linear(300, action_dim)  # 输出层，输出维度为动作维度
        self.max_action = max_action  # 动作的最大值，用于约束输出

    def forward(self, state):
        # 前向传播函数
        a = F.relu(self.fc1(state))  # 第一层激活函数ReLU
        a = F.relu(self.fc2(a))  # 第二层激活函数ReLU
        a = torch.tanh(self.fc3(a)) * self.max_action  # 输出层使用tanh激活并乘以最大动作值
        return a  # 返回动作值

# 定义Critic（值网络）类
class Critic(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(Critic, self).__init__()
        # 定义三层全连接层，输入是状态和动作的拼接，输出是Q值
        self.fc1 = nn.Linear(state_dim + action_dim, 400)  # 第一层全连接
        self.fc2 = nn.Linear(400, 300)  # 第二层全连接
        self.fc3 = nn.Linear(300, 1)  # 输出层，输出一个Q值

    def forward(self, state, action):
        # 将状态和动作拼接在一起作为输入
        state_action = torch.cat([state, action], 1)
        q = F.relu(self.fc1(state_action))  # 第一层激活函数ReLU
        q = F.relu(self.fc2(q))  # 第二层激活函数ReLU
        q = self.fc3(q)  # 输出层
        return q  # 返回Q值


# 定义TD3算法类
class TD3():
    def __init__(self, state_dim, action_dim, max_action):
        # 初始化Actor网络和目标Actor网络
        self.actor = Actor(state_dim, action_dim, max_action).to(device)
        self.actor_target = Actor(state_dim, action_dim, max_action).to(device)

        # 初始化两个Critic网络及其目标网络
        self.critic_1 = Critic(state_dim, action_dim).to(device)
        self.critic_1_target = Critic(state_dim, action_dim).to(device)
        self.critic_2 = Critic(state_dim, action_dim).to(device)
        self.critic_2_target = Critic(state_dim, action_dim).to(device)

        # 定义优化器
        self.actor_optimizer = optim.Adam(self.actor.parameters())  # Actor网络的优化器
        self.critic_1_optimizer = optim.Adam(self.critic_1.parameters())  # Critic 1网络的优化器
        self.critic_2_optimizer = optim.Adam(self.critic_2.parameters())  # Critic 2网络的优化器

        # 将目标网络的参数初始化为与主网络相同
        self.actor_target.load_state_dict(self.actor.state_dict())
        self.critic_1_target.load_state_dict(self.critic_1.state_dict())
        self.critic_2_target.load_state_dict(self.critic_2.state_dict())

        self.max_action = max_action  # 最大动作值
        self.memory = Replay_buffer(args.capacity)  # 初始化经验回放缓冲区
        self.writer = SummaryWriter(directory)  # 初始化TensorBoard记录器
        self.num_critic_update_iteration = 0  # Critic更新次数计数
        self.num_actor_update_iteration = 0  # Actor更新次数计数
        self.num_training = 0  # 总训练次数计数

    def select_action(self, state):
        # 根据当前策略选择动作
        state = torch.tensor(state.reshape(1, -1)).float().to(device)  # 将状态转换为张量
        return self.actor(state).cpu().data.numpy().flatten()  # 通过Actor网络生成动作

    def update(self, num_iteration):
        # 更新网络，训练过程
        if self.num_training % 500 == 0:  # 每500次训练打印日志
            print("====================================")
            print("model has been trained for {} times...".format(self.num_training))
            print("====================================")
        for i in range(num_iteration):  # 迭代更新
            x, y, u, r, d = self.memory.sample(args.batch_size)  # 从经验回放缓冲区中采样
            state = torch.FloatTensor(x).to(device)  # 转换采样的状态为张量
            action = torch.FloatTensor(u).to(device)  # 转换采样的动作为张量
            next_state = torch.FloatTensor(y).to(device)  # 转换采样的下一状态为张量
            done = torch.FloatTensor(d).to(device)  # 转换采样的完成标志为张量
            reward = torch.FloatTensor(r).to(device)  # 转换采样的奖励为张量

            # 选择目标动作，并加入噪声
            noise = torch.ones_like(action).data.normal_(0, args.policy_noise).to(device)
            noise = noise.clamp(-args.noise_clip, args.noise_clip)  # 裁剪噪声
            next_action = (self.actor_target(next_state) + noise).clamp(-self.max_action, self.max_action)

            # 计算目标Q值
            target_Q1 = self.critic_1_target(next_state, next_action)
            target_Q2 = self.critic_2_target(next_state, next_action)
            target_Q = torch.min(target_Q1, target_Q2)  # 取两个Critic网络中最小的Q值
            target_Q = reward + ((1 - done) * args.gamma * target_Q).detach()  # Bellman公式计算目标值

            # 优化Critic 1网络
            current_Q1 = self.critic_1(state, action)  # 当前Q值
            loss_Q1 = F.mse_loss(current_Q1, target_Q)  # 均方误差损失
            self.critic_1_optimizer.zero_grad()  # 清除梯度
            loss_Q1.backward()  # 反向传播
            self.critic_1_optimizer.step()  # 更新参数

            # 优化Critic 2网络
            current_Q2 = self.critic_2(state, action)
            loss_Q2 = F.mse_loss(current_Q2, target_Q)
            self.critic_2_optimizer.zero_grad()
            loss_Q2.backward()
            self.critic_2_optimizer.step()

            # 延迟更新策略网络和目标网络
            if i % args.policy_delay == 0:
                # 计算Actor损失
                actor_loss = - self.critic_1(state, self.actor(state)).mean()  # 最大化Q值
                self.actor_optimizer.zero_grad()  # 清除梯度
                actor_loss.backward()  # 反向传播
                self.actor_optimizer.step()  # 更新Actor网络

                # 更新目标网络参数
                for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()):
                    target_param.data.copy_((1 - args.tau) * target_param.data + args.tau * param.data)
                for param, target_param in zip(self.critic_1.parameters(), self.critic_1_target.parameters()):
                    target_param.data.copy_((1 - args.tau) * target_param.data + args.tau * param.data)
                for param, target_param in zip(self.critic_2.parameters(), self.critic_2_target.parameters()):
                    target_param.data.copy_((1 - args.tau) * target_param.data + args.tau * param.data)

                self.num_actor_update_iteration += 1  # 更新计数
        self.num_critic_update_iteration += 1  # Critic更新计数
        self.num_training += 1  # 总训练次数计数

    def save(self):
        # 保存模型的参数到指定目录
        torch.save(self.actor.state_dict(), directory + 'actor.pth')  # 保存Actor网络
        torch.save(self.actor_target.state_dict(), directory + 'actor_target.pth')  # 保存目标Actor网络
        torch.save(self.critic_1.state_dict(), directory + 'critic_1.pth')  # 保存Critic 1网络
        torch.save(self.critic_1_target.state_dict(), directory + 'critic_1_target.pth')  # 保存目标Critic 1网络
        torch.save(self.critic_2.state_dict(), directory + 'critic_2.pth')  # 保存Critic 2网络
        torch.save(self.critic_2_target.state_dict(), directory + 'critic_2_target.pth')  # 保存目标Critic 2网络
        print("====================================")
        print("Model has been saved...")  # 打印保存完成日志
        print("====================================")

    def load(self):
        # 加载保存的模型参数
        self.actor.load_state_dict(torch.load(directory + 'actor.pth'))  # 加载Actor网络
        self.actor_target.load_state_dict(torch.load(directory + 'actor_target.pth'))  # 加载目标Actor网络
        self.critic_1.load_state_dict(torch.load(directory + 'critic_1.pth'))  # 加载Critic 1网络
        self.critic_1_target.load_state_dict(torch.load(directory + 'critic_1_target.pth'))  # 加载目标Critic 1网络
        self.critic_2.load_state_dict(torch.load(directory + 'critic_2.pth'))  # 加载Critic 2网络
        self.critic_2_target.load_state_dict(torch.load(directory + 'critic_2_target.pth'))  # 加载目标Critic 2网络
        print("====================================")
        print("Model has been loaded...")  # 打印加载完成日志
        print("====================================")


# 主程序入口
if __name__ == '__main__':
    # 初始化TD3智能体
    agent = TD3(state_dim, action_dim, max_action)
    ep_r = 0  # 累计奖励初始化为0

    if args.mode == 'test':  # 如果模式为测试
        agent.load()  # 加载模型
        for i in range(args.iteration):  # 测试运行指定迭代次数
            state = env.reset()  # 环境重置，获取初始状态
            for t in count():  # 无限循环，直到完成或达到限制
                action = agent.select_action(state)  # 使用智能体选择动作
                next_state, reward, done, info = env.step(np.float32(action))  # 执行动作，获取下一状态和奖励
                ep_r += reward  # 累计奖励
                env.render()  # 渲染环境（显示界面）
                if done or t == 2000:  # 如果完成或者达到最大步数
                    print("Ep_i \t{}, the ep_r is \t{:0.2f}, the step is \t{}".format(i, ep_r, t))  # 打印日志
                    break  # 结束当前测试
                state = next_state  # 更新当前状态

    elif args.mode == 'train':  # 如果模式为训练
        print("====================================")
        print("Collection Experience...")  # 收集经验
        print("====================================")
        if args.load:
            agent.load()  # 如果需要，加载模型
        for i in range(args.num_iteration):  # 训练运行指定迭代次数
            state = env.reset()  # 环境重置
            for t in range(2000):  # 每次训练的最大步数
                # 使用智能体选择动作，并加入探索噪声
                action = agent.select_action(state)
                action = action + np.random.normal(0, args.exploration_noise, size=env.action_space.shape[0])
                action = action.clip(env.action_space.low, env.action_space.high)  # 限制动作范围
                next_state, reward, done, info = env.step(action)  # 执行动作
                ep_r += reward  # 累计奖励

                # 如果需要渲染且达到渲染间隔，显示环境
                if args.render and i >= args.render_interval:
                    env.render()

                # 将当前状态、下一状态、动作、奖励、完成标志存入经验回放缓冲区
                agent.memory.push((state, next_state, action, reward, np.float(done)))

                # 打印内存大小日志
                if (i + 1) % 10 == 0:
                    print('Episode {},  The memory size is {} '.format(i, len(agent.memory.storage)))

                # 如果经验回放缓冲区已满，则更新模型
                if len(agent.memory.storage) >= args.capacity - 1:
                    agent.update(10)  # 每次更新10步

                state = next_state  # 更新当前状态
                if done or t == args.max_episode - 1:  # 如果完成或达到最大轮数
                    agent.writer.add_scalar('ep_r', ep_r, global_step=i)  # 记录奖励到日志
                    if i % args.print_log == 0:
                        print("Ep_i \t{}, the ep_r is \t{:0.2f}, the step is \t{}".format(i, ep_r, t))  # 打印奖励日志
                    ep_r = 0  # 重置累计奖励
                    break  # 跳出当前循环

            # 每隔一定间隔保存模型
            if i % args.log_interval == 0:
                agent.save()

    else:
        raise NameError("mode wrong!!!")  # 如果模式错误，抛出异常

